#include "Common.h"
#include "ReadMemory.h"
#include "Public.h"
#include "WriteMemory.h"
#include "Common7.h"
#include "Common10.h"
#include "Handle.h"
#include "FarCall.h"
#include "EnumCallback.h"
#include "RegObCallBack.h"

extern PVOID  g_obProcessHandle1;
extern PVOID  g_obProcessHandle2;

NTSTATUS TraCommMsg(stCommonDate* msg)
{
	
	
	if (msg->id == 0x123456)
	{
		//COMM_TEST, READMEMORY, WRITEMEMORY
		//DbgPrintEx(77, 0, "cmd:%X pid:%I64X\r\n",msg->cmd, msg->pid);
		switch (msg->cmd)
		{
			case COMM_TEST:
				DbgPrintEx(77, 0, "[wz]DRIVE_TEST\r\n");
				msg->reStatus = STATUS_SUCCESS;
				break;
			case READMEMORY:
				msg->reStatus = ReadMemory2(msg->pid, (PVOID)msg->srcAddr, (PVOID)msg->dstAddr, msg->size);
				break;
			case WRITEMEMORY:
				msg->reStatus = WriteMemory(msg->pid, (PVOID)msg->srcAddr, (PVOID)msg->dstAddr, msg->size);
				break;
			case GETMODULEBASE:
				
				msg->reStatus = GetProcessModuleAddress(msg->pid,msg->srcAddr, msg->dstAddr);
				break;
			case UNINSTALL:
				UninstallCommon();
				break;
			case PROTECTEDPROCESS:
				//DbgPrintEx(77, 0, "PROTECTEDPROCESS pid:%I64X\r\n", msg->pid);
				msg->reStatus = ProtectByDestoryPspCidTable(msg->pid);
				break;
			case REMOTECALL:
				msg->reStatus = RemoteCall(msg->pid,msg->srcAddr,msg->size);
				break;
			case OBREGPORTECTED:
				msg->reStatus = ObRegisterCallBacksInit(msg->srcAddr);
				break;
			case ENUMCALLBACK:
				 EnumAllCallBack();
				msg->reStatus = STATUS_SUCCESS;
				break;
			case UNSETORREGCALLBACK:
				msg->reStatus = UnCallBack(msg->srcAddr, msg->dstAddr);
				break;
		default:
			break;
		}
	}
	else
	{
		return STATUS_FAIL_CHECK;
	}
	return STATUS_SUCCESS;
}





NTSTATUS InstallCommon()
{
	RTL_OSVERSIONINFOW verInfo = { 0 };

	verInfo.dwOSVersionInfoSize = sizeof(RTL_OSVERSIONINFOW);

	RtlGetVersion(&verInfo);

	DbgPrintEx(77, 0, "[wz]InstallCommon OS ver :%d\r\n", verInfo.dwBuildNumber);

	if (verInfo.dwBuildNumber == 7600 || verInfo.dwBuildNumber == 7601)
	{		
			return InstallCommonForWIn7(TraCommMsg);
	}
	else
	{
			return InstallCommonForWIn10(TraCommMsg);
	}

}
NTSTATUS UninstallCommon()
{
	RTL_OSVERSIONINFOW verInfo = { 0 };

	verInfo.dwOSVersionInfoSize = sizeof(RTL_OSVERSIONINFOW);

	RtlGetVersion(&verInfo);

	DbgPrintEx(77, 0, "[wz]UninstallCommon OS ver :%d\r\n", verInfo.dwBuildNumber);

	ObRegisterUnload();
	

	if (verInfo.dwBuildNumber == 7600 || verInfo.dwBuildNumber == 7601)
	{
		return UninstallCommonForWIn7();
	}
	else
	{
		return UninstallCommonForWIn10();
	}
	
	

}